#include <sys/socket.h>
#include <netinet/in.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>

#include <iostream>
#include <cstring>
#include <string>

#define BUFFER_LEN      4096

class TcpServerSelect {
    enum IoMode { Blocking, NonBlocking };
public:
    TcpServerSelect() : _port(8080), _listenBlock(20), _listenfd(-1), 
                        _ioMode(Blocking), _finished(false) {}

    int     initializer();
    void    run();

    void    setMode(int mode) { _ioMode = mode; }
    void    setBlock(int num) { _listenBlock = num; }
    void    setPort(short port) { _port = port; }
    void    setFinished(bool finished) { _finished = finished; }

protected:
    bool    setIoMode(int fd, int mode);
    void    acceptConnect(int& maxfd, fd_set& rfds);
    void    recvData(int clientfd, fd_set& wfds, fd_set& rfds);
    void    sendData(int clientfd, fd_set& wfds, fd_set& rfds);

protected:
    short       _port;
    int         _listenBlock;
    int         _listenfd;
    int         _ioMode;
    bool        _finished;
    std::string _strData;
};

int 
TcpServerSelect::initializer()
{
    if (_listenfd > 0) {
        std::cout << "The listening port already exists. listen fd " << _listenfd << std::endl;
        return 0;
    }
    // 1. Create socket
    _listenfd = socket(AF_INET, SOCK_STREAM, 0);
    if (_listenfd == -1) {
        std::cout << "socket return " << errno << ", " << strerror(errno) << std::endl;
        return -1;
    }

    // 2. Set the port and bind it.
    sockaddr_in serverAddr;
    memset(&serverAddr, 0, sizeof(sockaddr_in));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.s_addr = htons(INADDR_ANY); // bind ip address.
    serverAddr.sin_port = htons(_port);  // bind port.
    if (bind(_listenfd, (sockaddr*)&serverAddr, sizeof(serverAddr)) == -1) {
        std::cout << "bind return " << errno << ", " << strerror(errno) << std::endl;
        return -2;
    }

    if (_ioMode == NonBlocking) {
        // set nonblock mode.
        setIoMode(_listenfd, O_NONBLOCK);
    }

    // 3. listening port.
    if (listen(_listenfd, _listenBlock) == -1) {
        std::cout << "listen return " << errno << ", " << strerror(errno) << std::endl;
        return -3;
    }
    std::cout << "server listening port " << _port << std::endl;
    return 0;
}

void    
TcpServerSelect::run()
{
    if (_listenfd < 0) {
        std::cout << "Initialization not completed." << std::endl;
        return;
    }

    fd_set wfds, rfds, exceptfds;
    FD_ZERO(&wfds);
    FD_ZERO(&exceptfds);
    FD_ZERO(&rfds);
    FD_SET(_listenfd, &rfds);

    struct timeval timeout;
    // 设置超时时间
    timeout.tv_sec = 5; // 5秒
    timeout.tv_usec = 0;

    // 小技巧：通过中间变量将判断位和修改位分开。
    fd_set wset, rset;

    // 遍历多少个fd
    int maxfd = _listenfd;
    while (!_finished) {
        // 4. select
        wset = wfds;
        rset = rfds;
        // 注意，这里的参数是使用的rset和wset，将判断位和修改位分开。
        int ret = select(maxfd + 1, &rset, &wset, &exceptfds, &timeout);
        if (ret == 0)
           continue;
        if (ret < 0) {
            std::cout << "select return error: " << errno << ". " << strerror(errno) << std::endl;
            continue;
        }
        if (FD_ISSET(_listenfd, &rset))
            acceptConnect(maxfd, rfds);
        for (int i = _listenfd + 1; i <= maxfd; ++i) {
            if (FD_ISSET(i, &rset))
                recvData(i, wfds, rfds);
            else if (FD_ISSET(i, &wset))
                sendData(i, wfds, rfds);
        }
    }
    close(_listenfd);
    _listenfd = -1;
}

bool 
TcpServerSelect::setIoMode(int fd, int mode)
{
    int flag = fcntl(fd, F_GETFL, 0);
    if (flag == -1) {
        std::cout << "fcntl get flags return " << errno << ", " << strerror(errno) << std::endl;
        return false;
    }
    flag |= mode;
    if (fcntl(fd, F_SETFL, flag) == -1) {
        std::cout << "fcntl set flags return " << errno << ", " << strerror(errno) << std::endl;
        return false;
    }
    return true;
}

void    
TcpServerSelect::acceptConnect(int& maxfd, fd_set& rfds)
{
    // 4. accept connect.
    sockaddr_in clientAddr;
    memset(&clientAddr, 0, sizeof(clientAddr));
    socklen_t clienLen = sizeof(clientAddr);
    int clientfd = accept(_listenfd, (sockaddr *)&clientAddr, &clienLen);
    if (clientfd == -1) {
        std::cout << "accept return " << errno << ", " << strerror(errno) << std::endl;
        return;
    }
    std::cout << "client fd " << clientfd << std::endl;

    FD_SET(clientfd, &rfds);
    if (clientfd > maxfd)
        maxfd = clientfd;
}

void    
TcpServerSelect::recvData(int clientfd, fd_set& wfds, fd_set& rfds)
{
    // 6. recv message
    char buffer[BUFFER_LEN];
    int ret = recv(clientfd, buffer, BUFFER_LEN, 0);
    if (ret == 0) {
        std::cout << "client " << clientfd << " connection dropped" << std::endl;
        close(clientfd);
        // clear read ready.
        FD_CLR(clientfd, &rfds);
        return;
    } else if (ret == -1) {
        std::cout << "recv buffer return " << errno << ", " << strerror(errno) << std::endl;
        return;
    }
    std::cout << "recv buffer from client "<< clientfd << ": " << buffer << std::endl;
    _strData = buffer;
    FD_CLR(clientfd, &rfds);
    // set fd send data ready.
    FD_SET(clientfd, &wfds);
}

void    
TcpServerSelect::sendData(int clientfd, fd_set& wfds, fd_set& rfds)
{
    if (_strData.empty())
        _strData = "Hello, Client!";
    // 5. send message.
    if (send(clientfd, _strData.c_str(), _strData.size(), 0) == -1) {
        std::cout << "send buffer return " << errno << ", " << strerror(errno) << std::endl;
        return;
    }
    // clear write fd
    FD_CLR(clientfd, &wfds);
    // set fd send data ready.
    FD_SET(clientfd, &rfds);
    
}

int main(int argc, char**argv)
{
    TcpServerSelect server;
    if (server.initializer() != 0)
        return -1;
    server.run();
}
